package cn.topkinson.es.etl.utils;

import cn.topkinson.es.etl.entity.es.*;
import com.alibaba.fastjson.JSON;
import org.elasticsearch.action.DocWriteResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

public class EsHelper {

    private final static Logger log = LoggerFactory.getLogger(EsHelper.class);

    @Autowired
    private RestHighLevelClient restHighLevelClient;

    /**
     * 搜索文档，返回指定类
     * @param searchDto
     * @param clazz
     * @param <T>
     * @return
     */
    public <T> SearchResult<T> pageHits(SearchDto searchDto, Class<T> clazz) {
        SearchResponse response = this.searchResponse(searchDto);
        if (response != null) {
            SearchHits hits = response.getHits();
            Long totalShards = hits.getTotalHits().value;
            List<T> list = new ArrayList<>();
            if (clazz != null) {
                for (SearchHit hit : hits) {
                    list.add(JSON.parseObject(hit.getSourceAsString(), clazz));
                }
            }
            return new SearchResult<>(totalShards, list, hits);
        }
        return null;
    }

    /**
     * 搜索文档
     * @param searchDto
     * @return
     */
    private SearchResponse searchResponse(SearchDto searchDto) {
        IndexEnum indexEnum = searchDto.getIndexEnum();
        if (indexEnum == null) {
            log.error("search with no IndexEnum");
            return null;
        }
        //请求对象
        SearchRequest searchRequest = new SearchRequest(searchDto.getIndexEnum().getIndex());
        //查询资源构造器
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        //返回页数的确切值
        sourceBuilder.trackTotalHits(true);
        //查询条件
        sourceBuilder.query(searchDto.getQueryBuilder());
        //分页设置
        Integer from = searchDto.getFrom();
        Integer size = searchDto.getSize();
        if (Objects.nonNull(from) && Objects.nonNull(size)) {
            sourceBuilder.from(searchDto.getFrom());
            sourceBuilder.size(searchDto.getSize());
        }
        //排序设置
        Map<String, SortOrder> sortOrderMap = searchDto.getSortOrderMap();
        if (!CollectionUtils.isEmpty(sortOrderMap)) {
            sortOrderMap.forEach((k, v) -> {
                sourceBuilder.sort(k, v);
            });
        }
        //结果集合
        sourceBuilder.fetchSource(searchDto.getIncludes(), searchDto.getExcludes());
        //将builder作为source设置到请求中
        searchRequest.source(sourceBuilder);
        try {
            log.info(String.format("ES search index【%s】 DSL : %s", indexEnum.getIndex(), sourceBuilder.toString()));
            long begin = System.currentTimeMillis();
            //发送请求，获得response，返回结果
            SearchResponse response = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info(String.format("ES search index【%s】 spend time : %s ms", indexEnum.getIndex(),
                    System.currentTimeMillis() - begin));
            log.info(String.format("ES search index【%s】 total hits : %s ", indexEnum.getIndex(),
                    response.getHits().getTotalHits().value));
            return response;
        } catch (Exception e) {
            log.error("search response exception", e);
        }
        return null;
    }

    /**
     * 更新文档，若不存在，则创建新文档
     * @param dto
     * @return
     * @throws IOException
     */
    public boolean upsert(UpdateDto dto) throws IOException {
        // 获取请求参数
        String id = dto.getId();
        String index = dto.getIndexEnum().getVersion();
        Map<String, Object> data = dto.getData();
        // 更新创建
        boolean rs = false;
        IndexRequest indexRequest = new IndexRequest(index).id(id).source(data, XContentType.JSON);
        UpdateRequest updateRequest = new UpdateRequest(index, id).doc(data, XContentType.JSON).upsert(indexRequest);
        UpdateResponse response = restHighLevelClient.update(updateRequest, RequestOptions.DEFAULT);
        if (response.getResult() == DocWriteResponse.Result.CREATED
                || response.getResult() == DocWriteResponse.Result.UPDATED
                || response.getResult() == DocWriteResponse.Result.NOOP) {
            rs = true;
        }
        log.info(String.format("ES upsert, index【%s】 result【%s】", index, response.getResult()));
        return rs;
    }

    /**
     * 删除索引数据
     * @param dto
     * @return
     * @throws IOException
     */
    public boolean delete(DeleteDto dto) throws IOException {
        // 获取请求参数
        String id = dto.getId();
        String index = dto.getIndexEnum().getVersion();
        // 删除
        boolean rs = false;
        DeleteRequest deleteRequest = new DeleteRequest(index, id);
        DeleteResponse response = restHighLevelClient.delete(deleteRequest, RequestOptions.DEFAULT);
        if (response.getResult() == DocWriteResponse.Result.DELETED
                || response.getResult() == DocWriteResponse.Result.NOT_FOUND) {
            rs = true;
        }
        log.info(String.format("ES delete, index【%s】 result【%s】", index, response.getResult()));
        return rs;
    }

}
